import argparse
import torch
import torch.nn as nn
import torch.nn.parallel
import torch.backends.cudnn as cudnn
import torch.optim
import torch.utils.data
import torch.utils.data.distributed
from train import *
from save import *
    

    
def print_args(args):
    print("\nArguments:")
    print("-" * 30)
    for key, value in vars(args).items():
        print(f"{key: <15}: {value}")
    print("-" * 30)

def parser_args():
    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
    parser.add_argument('--data_path', default='/home/course/sea_with_stars/data', type=str)
    parser.add_argument('-n', '--num_class', default=10, type=int, metavar='N',)
    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',)
    parser.add_argument('-b', '--batch-size', default=64, type=int,metavar='N',)
    parser.add_argument('--weight', default='/home/course/sea_with_stars/weight/best_model.pth', type=str,)
    parser.add_argument('--gpu', default=0, type=int,)
    parser.add_argument('-p', '--print-freq', default=10, type=int, metavar='N', )
    args = parser.parse_args()
    print_args(args)
    return args


def main(args):

    device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    model = get_model(args.num_class)
    model.apply(inplace_relu)
    model = model.cuda(args.gpu)
    state_dict = torch.load(args.weight, map_location=device)
    model.load_state_dict(state_dict)
    test_dataset = PointCloudDataset(root=args.data_path, split='test')
    val_loader = DataLoader(test_dataset, batch_size=1, shuffle=False,collate_fn=pad_collate_fn)
    criterion = nn.CrossEntropyLoss().cuda(args.gpu)
    cudnn.benchmark = True
    model.eval()
    with torch.no_grad():
        # for i, (points, target) in enumerate(val_loader):
        #     points,target =  points.cuda(args.gpu), target.cuda(args.gpu)
            
            points=load_txt("/home/course/sea_with_stars/data/number7.txt.txt")
            points = points.transpose(2, 1)
            pred, _ = model(points)
            pred_choice = pred.data.max(1)[1]
            print("pred_choice is ",pred_choice,"target is ",7)
            return 0

if __name__ == '__main__':
    args=parser_args()
    main(args)